# Copyright (c) 2018 Presto Labs Pte. Ltd.
# Author: yuxuan

import datetime

import numpy as np
from absl import app
from absl.flags import FLAGS, DEFINE_string

from coin.database.builder.cache_builder_base2 import CacheBuilderBase2
from coin.strategy.mm.dumper_base import run_from_archive
from coin.support.feed_tool.feed_research.all_products import (
    gen_all_products,
    gen_all_feed_sub_request,
)
from coin.strategy.mm.subscription import FeedSubscriptionRequest
from coin.database.builder.archive_feed_cache_builder import ArchiveFeedCacheBuilder

from coin.database.builder.bitmex_archive_feed_cache_builder import BitmexArchiveFeedCacheBuilder
from coin.database.builder.okex_archive_feed_cache_builder import OkexArchiveFeedCacheBuilder
from coin.database.builder.okex_v3_archive_feed_cache_builder import (
    OkexV3FuturesArchiveFeedCacheBuilder,
    OkexV3SwapArchiveFeedCacheBuilder,
)

DEFINE_string('machine', '', '')

DEFINE_string('feed_root', None, '')


def gen_all_products_symbols(feed_sub_request):
  products = []
  full_symbols = []
  for sub_request, _products in feed_sub_request.gen_iter():
    products.extend(_products)
    if sub_request.api_version in ['v3', 'v3_swap']:
      full_symbols.extend([product.full_symbol_v3 for product in _products])
    else:
      full_symbols.extend([product.full_symbol for product in _products])
  return products, full_symbols


class CacheBuilder(CacheBuilderBase2):
  def __init__(self, date):
    CacheBuilderBase2.__init__(self, date)
    self._date = datetime.datetime.strptime(str(date), '%Y%m%d')
    self._date_int = date
    self._feed_sub_request = gen_all_feed_sub_request(self._date)

  def _get_index_attributes(self):
    return ['base', 'quote', 'exchange', 'symbol', 'expiry']

  def build_index(self):
    trading_date = self._date
    symbols_idx = dict()
    products, symbols = gen_all_products_symbols(self._feed_sub_request)
    assert len(products) == len(symbols)
    for symbol in symbols:
      assert symbol not in symbols_idx, (symbol, symbols_idx)
      symbols_idx[symbol] = len(symbols_idx)
    all_dates_int = [int(trading_date.strftime('%Y%m%d'))]
    self.build_cache('metadata',
                     'presto',
                     'dates',
                     np.array(all_dates_int),
                     True, {
                         'align': 'd', 'type': 'vector'
                     })
    self.build_cache('metadata',
                     'presto',
                     'symbols',
                     np.array(symbols),
                     True, {
                         'align': 'i', 'type': 'vector'
                     })
    for attribute in self._get_index_attributes():
      data = [['' for _ in range(len(symbols))] for _ in range(len(all_dates_int))]
      for di, trading_date in enumerate([trading_date]):
        for _i, symbol in enumerate(symbols):
          product = products[_i]
          ii = symbols_idx[symbol]
          if hasattr(product, attribute):
            data[di][ii] = getattr(product, attribute)
      self.build_cache('metadata',
                       'presto',
                       attribute,
                       np.array(data).astype(str),
                       True, {
                           'align': 'di', 'type': 'matrix'
                       })

  def get_symbols_idx_map(self):
    self.ensure_initialized()
    return {symbol: i for i, symbol in enumerate(self._all_symbols)}

  @staticmethod
  def _get_interval_end_points(start, end, interval):
    start = int(start)
    end = int(end)
    interval = int(interval)
    seps = [i for i in range(start + interval, end + interval, interval)]
    assert end == seps[-1]
    return seps

  def build_open_interest(self, interval, interval_name):
    columns = ['open_interest']
    symbols_idx_map = self.get_symbols_idx_map()
    start = self._date
    end = start + datetime.timedelta(days=1)
    interval_ends = self._get_interval_end_points(start.timestamp() * 10**9,
                                                  end.timestamp() * 10**9,
                                                  interval)
    all_intervals = [start.timestamp() * 10**9] + interval_ends

    data = {
        column: {
            price_type: np.full((len(interval_ends), len(self._all_symbols)), np.nan)
            for price_type in ArchiveFeedCacheBuilder.price_types()
        } for column in columns
    }

    for cls in [
        BitmexArchiveFeedCacheBuilder,
        OkexV3FuturesArchiveFeedCacheBuilder,
        OkexV3SwapArchiveFeedCacheBuilder
    ]:
      products = self._feed_sub_request.get_products(cls._sub_request)
      feed_sub_request = FeedSubscriptionRequest.create(products, cls._sub_request)
      archive_builder = cls(products, columns, interval_ends)
      run_from_archive(feed_sub_request,
                       archive_builder.on_book_reset,
                       start,
                       end,
                       machine=[FLAGS.machine],
                       root_dir=FLAGS.feed_root)
      for column in columns:
        for product in products:
          if cls._sub_request.api_version in ['v3', 'v3_swap']:
            full_symbol = product.full_symbol_v3
          else:
            full_symbol = product.full_symbol
          iidx = symbols_idx_map[full_symbol]
          for price_type in ArchiveFeedCacheBuilder.price_types():
            data[column][price_type][:, iidx] = \
                archive_builder.get_data(product, column)[price_type]

    for column, column_data in data.items():
      self.build_cache(column,
                       'futures',
                       'intervals',
                       np.array(all_intervals[:-1]),
                       True, {
                           'align': 'i', 'interval': interval_name
                       })
      for price_type, a_data in column_data.items():
        self.build_cache(column,
                         'futures',
                         '%s_%s' % (price_type, column),
                         a_data,
                         True, {
                             'interval': interval_name, 'align': 'ti', 'type': 'matrix'
                         })


def build_cache_batch(date):
  builder = CacheBuilder(date)
  builder.build_index()
  builder.build_open_interest(10 * 60 * 10**9, '10min')


def main(args):
  # Uncomment below to build cache
  # build_cache_batch(20181101)
  return


if __name__ == '__main__':
  app.run(main)
